// hello.cpp
#include <pybind11/numpy.h>
#include <stdexcept>

namespace py = pybind11;

py::array_t<float> matmul(const py::array_t<float>& a, const py::array_t<float>& b) {
    if (a.ndim() != 2 || b.ndim() != 2)
        throw std::runtime_error("輸入必須是兩個 2D 矩陣");

    auto a_buf = a.unchecked<2>();
    auto b_buf = b.unchecked<2>();

    ssize_t M = a_buf.shape(0);
    ssize_t K = a_buf.shape(1);
    ssize_t N = b_buf.shape(1);

    if (b_buf.shape(0) != K)
        throw std::runtime_error("維度不一致：A的列數必須等於B的行數");

    auto result = py::array_t<float>({M, N});
    auto r = result.mutable_unchecked<2>();

    for (ssize_t i = 0; i < M; i++) {
        for (ssize_t j = 0; j < N; j++) {
            float sum = 0.0;
            for (ssize_t k = 0; k < K; k++) {
                sum += a_buf(i, k) * b_buf(k, j);
            }
            r(i, j) = sum;
        }
    }

    return result;
}